open Unix

module FD =
struct 
  type t = file_descr
  let compare x y = compare (Hashtbl.hash x) (Hashtbl.hash y)
end

module FDS = Set.Make(FD)
module H = Hashtbl

let ht = H.create 17

exception Disconnected

let read_from so =
  let len = 100 in
  let smbuf = String.create len in
  let buf = Buffer.create 100 in
  let rec read_till () =
    let i = recv so smbuf 0 len [] in
      if i = 0 then raise Disconnected ;
      if i < len then
        Buffer.add_string buf (String.sub smbuf 0 i) 
      else
        let () = Buffer.add_string buf smbuf in
          read_till ()
  in
    try 
      read_till () ; buf
    with
      | Unix_error (EAGAIN,s1,s2) ->  buf
      
let write_to so s = 
  ignore (send so (Buffer.contents s) 0 (Buffer.length s) [])

let remove_error_socks errors ss =
  List.fold_left
    (fun acc x -> FDS.remove x acc) ss errors

let handle_incoming ready sock =
  let rec aux ns bs = function
    | [] -> ns, bs
    | (so::sos) ->
        if so = sock then
          let so, _ = accept sock in begin
            set_nonblock so;
            H.add ht so (Queue.create ());
            aux (FDS.add so ns) bs sos
          end
        else
          let bs = 
            try
              let s = read_from so in
                H.iter (fun soc q -> 
                          if soc = so then () else Queue.add s q) ht;
                bs
            with
              | Disconnected -> FDS.add so bs
          in
            aux ns bs sos
  in
    aux FDS.empty FDS.empty ready

let handle_outgoing write sock =
  let rec waux q so =
    let b = Queue.pop q in
      write_to so b;
      waux q so
  in
  let rec aux bs = function
    | [] -> bs
    | (so::sos) ->
        if so = sock then aux bs sos
        else
          try 
            waux (H.find ht so) so
          with
            | Not_found | Queue.Empty -> aux bs sos
  in
    aux FDS.empty write



let select_loop sock =
  let rec aux ss =
      let sl = FDS.elements ss in
      Format.printf "having %d sockets in contact ...@. " (List.length sl) ;
      let sls = sock::sl in
      let read,write,errors = select sls [] sls (-1.0) in
     (* remove all socks from list that provoke an error; probably a bit cruel
      * *)
        Format.printf "new data for read ...@. " ;
        let ss = remove_error_socks errors ss in
        let ss = 
          match read with
            | [] -> ss (* no ready socks, go on *)
            | l -> 
                (* handle client connections and new read data *)
                (* get list of new socks and bad socks *)
                let ns, bs = handle_incoming l sock in
                (* throw away disc socks and add new socks *)
                let ss = FDS.union (FDS.diff ss bs) ns in
                let () = FDS.iter (H.remove ht) bs in
                  ss
        in
          Format.printf "checking for write ...@. " ;
          let sl = FDS.elements ss in
          let _,write,errors = select [] sl sl 0.1 in
          let ss = remove_error_socks errors ss in
            Format.printf "handling writes ...@. " ;
            let ss = 
              match write with
                | [] -> ss
                | l ->
                    let bs = handle_outgoing l sock in
                    let ss = FDS.diff ss bs in
                    let () = FDS.iter (H.remove ht) bs in
                      ss
            in
              aux ss
  in
    aux FDS.empty

let _ =
  let lb = ADDR_INET (inet_addr_loopback, 5000) in
  let sock = socket PF_INET SOCK_STREAM 0 in
  let () = set_nonblock sock in
    try
      Format.printf "binding to port ...@.";
      bind sock lb ;
      Format.printf "listening to port ...@.";
      listen sock 3;
      Format.printf "entering select loop ...@.";
      select_loop sock
    with Unix.Unix_error (e,_,_) ->
      Printf.printf "%s\n" (error_message e);
      H.iter (fun s _ -> shutdown s SHUTDOWN_ALL) ht;
      shutdown sock SHUTDOWN_ALL
